// vectordatabase.cpp - 文件存储版本
#include "vectordatabase.h"
#include <QDebug>
#include <QtMath>
#include <QJsonDocument>
#include <QStandardPaths>

VectorDatabase::VectorDatabase(QObject *parent)
    : QObject(parent), m_nextId(1)
{
}

bool VectorDatabase::insertDocument(const DocumentChunk &chunk)
{
    DocumentChunk newChunk = chunk;
    newChunk.id = m_nextId++;
    m_documents.append(newChunk);
    // 每次插入后保存
    qDebug() << "Inserted document with id:" << newChunk.id;
    return saveToFile();
}

QVector<DocumentChunk> VectorDatabase::searchSimilar(const QVector<float> &queryVector, int topK)
{
    QVector<DocumentChunk> results;
    QVector<QPair<float, DocumentChunk>> candidates;
    qDebug() << "Searching in" << m_documents.size() << "documents";

    for (const DocumentChunk &chunk : m_documents) {
        float similarity = cosineSimilarity(queryVector, chunk.embedding);
        candidates.append(qMakePair(similarity, chunk));
    }

    // 按相似度排序
    std::sort(candidates.begin(), candidates.end(),
              [](const QPair<float, DocumentChunk> &a, const QPair<float, DocumentChunk> &b) {
                  return a.first > b.first;
              });

    // 返回前topK个结果
    for (int i = 0; i < qMin(topK, candidates.size()); ++i) {
        results.append(candidates[i].second);
    }
    return results;
}

bool VectorDatabase::saveToFile()
{
//    qDebug() << "saveToFile";
    QJsonDocument doc;
    QJsonObject root;
    QJsonArray documentsArray;
    QString dirPath = QDir::currentPath() + "/KyAiData";
    QDir dir(dirPath);
    if (!dir.exists()) {
        if (!dir.mkpath(".")) {
            qDebug() << "创建目录失败" << dirPath;
        }
    }
    QFile file(m_filePath);

    for (const DocumentChunk &chunk : m_documents) {
        documentsArray.append(chunk.toJson());
    }

    root["documents"] = documentsArray;
    root["next_id"] = m_nextId;

    doc.setObject(root);
    if (!file.open(QIODevice::WriteOnly)) {
        qDebug() << "Failed to open file for writing";
        return false;
    }

    file.write(doc.toJson());
    file.close();

    return true;
}

bool VectorDatabase::loadFromFile()
{
    qDebug() << "loadFromFile";
    if (!QFile::exists(m_filePath)) {
        QFile file(m_filePath);
        if (file.open(QIODevice::WriteOnly)) {
            file.write("[]"); // 空JSON数组
            file.close();
            m_documents.clear();  // 新数据库
            m_nextId = 1;
        }
        qDebug() << "Database file doesn't exist, starting fresh";
    }
    QFile file(m_filePath);
    if (!file.open(QIODevice::ReadOnly)) {
        qDebug() << "Failed to open file for reading" ;
        return false;
    }

    QByteArray data = file.readAll();
//    qDebug() << "File content:" << data;
    file.close();

    QJsonParseError error;
    QJsonDocument doc = QJsonDocument::fromJson(data, &error);

    if (error.error != QJsonParseError::NoError) {
        qDebug() << "JSON parse error:" << error.errorString();
        return false;
    }
    QJsonObject root = doc.object();
    QJsonArray documentsArray = root["documents"].toArray();
    m_nextId = root["next_id"].toInt(1);

    m_documents.clear();
    for (const QJsonValue &value : documentsArray) {
        DocumentChunk chunk;
        chunk.fromJson(value.toObject());
        m_documents.append(chunk);
    }

    qDebug() << "Loaded" << m_documents.size() << "documents from file";
    return true;
}

float VectorDatabase::cosineSimilarity(const QVector<float> &vec1, const QVector<float> &vec2)
{
    if (vec1.size() != vec2.size()) {
        return 0.0f;
    }

    float dotProduct = 0.0f;
    float norm1 = 0.0f;
    float norm2 = 0.0f;

    for (int i = 0; i < vec1.size(); ++i) {
        dotProduct += vec1[i] * vec2[i];
        norm1 += vec1[i] * vec1[i];
        norm2 += vec2[i] * vec2[i];
    }

    if (norm1 == 0.0f || norm2 == 0.0f) {
        return 0.0f;
    }

    return dotProduct / (qSqrt(norm1) * qSqrt(norm2));
}

bool VectorDatabase::documentExists(const QString &filePath, const QString &text) const
{
    for (const DocumentChunk &doc : m_documents) {
        if (doc.filePath == filePath && doc.text == text) {
            return true;
        }
    }
    return false;
}

// DocumentChunk的序列化方法
QJsonObject DocumentChunk::toJson() const
{
//    qDebug() << "toJson";
    QJsonObject obj;
    obj["id"] = id;
    obj["text"] = text;
    obj["file_path"] = filePath;

    QJsonArray embeddingArray;
    for (float value : embedding) {
        embeddingArray.append(static_cast<double>(value));
    }
    obj["embedding"] = embeddingArray;

    return obj;
}

void DocumentChunk::fromJson(const QJsonObject &json)
{
//    qDebug() << "fromJson";
    id = json["id"].toInt();
    text = json["text"].toString();
    filePath = json["file_path"].toString();

    embedding.clear();
    QJsonArray embeddingArray = json["embedding"].toArray();
    for (const QJsonValue &value : embeddingArray) {
        embedding.append(static_cast<float>(value.toDouble()));
    }
}

bool VectorDatabase::clearDatabase()
{
    // 清空内存中的数据
    m_documents.clear();
    m_nextId = 1;

    // 删除磁盘上的文件
    QFile file(m_filePath);
    if (file.exists()) {
        if (!file.remove()) {
            qDebug() << "Failed to remove database file";
            return false;
        }
    }

    // 创建一个新的空数据库文件
    if (!file.open(QIODevice::WriteOnly)) {
        qDebug() << "Failed to create new database file";
        return false;
    }

    file.write("{\"documents\":[],\"next_id\":1}");
    file.close();

    qDebug() << "Database cleared successfully";
    return true;
}

